﻿/*
    Copyright (c) 2012, Matías López
    All rights reserved.

    Redistribution and use in source and binary forms, with or without modification, are permitted 
    provided that the following conditions are met:
        *   Redistributions of source code must retain the above copyright notice, this list of 
            conditions and the following disclaimer.
        *   Redistributions in binary form must reproduce the above copyright notice, this list of 
            conditions and the following disclaimer in the documentation and/or other materials 
            provided with the distribution.
        *   Neither the name of Matías López nor the names of its contributors may be used to 
            endorse or promote products derived from this software without specific prior written 
            permission.
    
    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR 
    IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY 
    AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 
    CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
    CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 
    SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 
    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR 
    OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
    POSSIBILITY OF SUCH DAMAGE.
 */

using System;
using System.Collections.Generic;
using System.Linq;
using System.Linq.Expressions;
using System.Text;
using System.Reflection;
using System.Data.Services.Client;

using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.StorageClient;
using System.Collections.ObjectModel;

namespace mana.Data.Azure.Tables
{
    /// <summary>
    /// Exposes all kinds of entities operations to execute against an Azure Table.
    /// </summary>
    /// <typeparam name="T">Type of entities that this context will handle.</typeparam>
    public class GenericTableContext<T> where T : TableServiceEntity
    {
        #region Delegates

        delegate void ContextOperation(T entity, bool skipExistenceCheck);

        #endregion

        #region Events

        /// <summary>
        /// Raised when any change is saved in the context. Especially useful during batch operations.
        /// </summary>
        public event EventHandler<GenericTableContextOperationEventArgs> ChangesSaved;

        #endregion

        private Type _entityType;
        private string _entitySetName;
        private PropertyInfo[] _entityProperties;
        private ConstructorInfo _entityConstructor;

        private bool _saveWithRetries;
        private int _batchSize;
        
        private TableServiceContext _context;

        private GenericTableContextOperationEventArgs _currentOperationEventArgs;

        /// <summary>
        /// Initializes a new instance of GenericTableContext targeting a Cloud Storage Account.
        /// </summary>
        /// <param name="account">Target Cloud Storage Account.</param>
        public GenericTableContext(CloudStorageAccount account)
            : this(account, null)
        { }

        /// <summary>
        /// Initializes a new instance of GenericTableContext targeting a Cloud Storage Account and with a specified Azure Table name.
        /// </summary>
        /// <param name="account">Target Cloud Storage Account.</param>
        /// <param name="entitySetName">Azure table name.</param>
        public GenericTableContext(CloudStorageAccount account, string entitySetName)
        {
            _entityType = typeof(T);
            
            if (string.IsNullOrWhiteSpace(entitySetName))
                entitySetName = _entityType.Name;
            _entitySetName = entitySetName;
            _entityProperties = _entityType.GetProperties();
            _entityConstructor = _entityType.GetConstructor(new Type[0] { });

            _batchSize = 100;
            
            _context = account.CreateCloudTableClient().GetDataServiceContext();
            _context.IgnoreResourceNotFoundException = true;
            _context.SaveChangesDefaultOptions = SaveChangesOptions.ReplaceOnUpdate;
            _context.MergeOption = MergeOption.PreserveChanges;
            _context.ResolveType = this.ResolveType;
            
            RetryPolicy = RetryPolicies.NoRetry();

            _currentOperationEventArgs = new GenericTableContextOperationEventArgs();
        }

        #region Properties

        /// <summary>
        /// Azure Table name.
        /// </summary>
        public string EntitySetName
        {
            get { return _entitySetName; }
        }

        /// <summary>
        /// Represents the entities in a Azure Table in a query. All queries to this context MUST target this property.
        /// </summary>
        public DataServiceQuery<T> Entities
        {
            get { return _context.CreateQuery<T>(_entitySetName); }
        }

        public TableServiceContext Context
        {
            get { return _context; }
        }

        /// <summary>
        /// Gets or sets the number of entities sent in each batch call.
        /// </summary>
        public int BatchSize
        {
            get { return _batchSize; }
            set { _batchSize = value; }
        }

        /// <summary>
        /// Gets or sets the Retry Policy.
        /// </summary>
        public RetryPolicy RetryPolicy
        {
            get { return _context.RetryPolicy; }
            set 
            { 
                _context.RetryPolicy = value;
                _saveWithRetries = !value.Equals(RetryPolicies.NoRetry());
            }
        }

        #endregion Properties

        #region Methods

        #region CRUD Operations

        #region CREATE

        /// <summary>
        /// Adds the specified entity to the Azure Table.
        /// </summary>
        /// <param name="entity"></param>
        public void Add(T entity)
        {
            ApplyOperation(entity, AddOperation);
        }

        /// <summary>
        /// Adds the specified entities to the Azure Table in batch. It groups entities by PartitionKey acording to the BatchSize specified.
        /// </summary>
        /// <param name="entities"></param>
        public void AddInBatch(IEnumerable<T> entities)
        {
            ApplyOperationInBatch(entities.ToList(), new ContextOperation(AddOperation));
        }

        #endregion CREATE

        #region RETRIEVE

        /// <summary>
        /// Gets an entity with the specified PartitionKey and RowKey.
        /// </summary>
        /// <param name="partitionKey">Entity's Partition Key.</param>
        /// <param name="rowKey">Entity's Row Key.</param>
        /// <returns>Returns the entity with the specified PartitionKey and RowKey.</returns>
        public T Get(string partitionKey, string rowKey)
        {
            var query = 
                from    entity in Entities
                where   entity.PartitionKey == partitionKey && entity.RowKey == rowKey
                select  entity;

            return ExecuteQuery(query).FirstOrDefault();
        }

        /// <summary>
        /// Gets the results of the specified query.
        /// </summary>
        /// <typeparam name="TResult">Result type that the query returns.</typeparam>
        /// <param name="query">Query to execute.</param>
        /// <returns>Returns the query results.</returns>
        public IEnumerable<TResult> Get<TResult>(IQueryable<TResult> query)
        {
            return ExecuteQuery(query);
        }

        /// <summary>
        /// Gets the first result of the specified query.
        /// </summary>
        /// <typeparam name="TResult">Result type that the query returns.</typeparam>
        /// <param name="query">Query to execute.</param>
        /// <returns>Returns the first query result.</returns>
        public TResult GetFirst<TResult>(IQueryable<TResult> query)
        {
            return ExecuteQuery(query).FirstOrDefault();
        }

        /// <summary>
        /// Gets the single expected result of the specified query.
        /// </summary>
        /// <typeparam name="TResult">Result type that the query returns.</typeparam>
        /// <param name="query">Query to execute.</param>
        /// <returns>Returns the single query expected result.</returns>
        public TResult GetSingle<TResult>(IQueryable<TResult> query)
        {
            return ExecuteQuery(query).SingleOrDefault();
        }

        /// <summary>
        /// Gets the entities of the specified PartitionKey.
        /// </summary>
        /// <param name="partitionKey">Entity's Partition Key.</param>
        /// <returns>Returns the entities with the specified PartitionKey.</returns>
        public IEnumerable<T> GetByPartitionKey(string partitionKey)
        {
            var query =
                from    entity in Entities
                where   entity.PartitionKey == partitionKey
                select  entity;

            return ExecuteQuery(query);
        }

        /// <summary>
        /// Gets all entities in this Azure Table.
        /// </summary>
        /// <returns>Returns all entities.</returns>
        public IEnumerable<T> GetAll()
        {
            return Get(Entities);
        }

        private IEnumerable<TResult> ExecuteQuery<TResult>(IQueryable<TResult> query)
        {
            return query.AsTableServiceQuery().Execute();
        }

        #endregion RETRIEVE

        #region UPDATE

        /// <summary>
        /// Updates the specified entity.
        /// </summary>
        /// <param name="entity">Entity to update.</param>
        public void Update(T entity)
        {
            Update(entity, false);
        }

        /// <summary>
        /// Updates the specified entity. Specifies whether to skip the existence check for the specified entity.
        /// </summary>
        /// <param name="entity">Entity to update.</param>
        /// <param name="skipExistenceCheck">Specifies to skip the entity's existence check or not.</param>
        public void Update(T entity, bool skipExistenceCheck)
        {
            ApplyOperation(entity, UpdateOperation, skipExistenceCheck);
        }

        #endregion UPDATE

        #region DELETE

        /// <summary>
        /// Deletes the entity with the specified PartitionKey 
        /// </summary>
        /// <param name="partitionKey">Entity's PartitionKey to delete.</param>
        /// <param name="rowKey">Entity's RowKey to delete</param>
        public void Delete(string partitionKey, string rowKey)
        {
            Delete(CreateInstanceOfEntity(partitionKey, rowKey));
        }

        /// <summary>
        /// Deletes the specified entity.
        /// </summary>
        /// <param name="entity">Entity to delete.</param>
        public void Delete(T entity)
        {
            ApplyOperation(entity, DeleteOperation, true);
        }

        /// <summary>
        /// Deletes the resulting entities of the query execution in batch.
        /// </summary>
        /// <param name="query">Query to execute.</param>
        public void DeleteInBatch(IQueryable<T> query)
        {
            var newQuery =
                from en in query
                select new { PartitionKey = en.PartitionKey, RowKey = en.RowKey };

            var q1 = this.ExecuteQuery(newQuery);

            List<T> entities = new List<T>();
            foreach (var entity in q1)
            {
                entities.Add(CreateInstanceOfEntity(entity.PartitionKey, entity.RowKey));
            }

            DeleteInBatch(entities);
        }

        /// <summary>
        /// Deletes the specified entities in batch.
        /// </summary>
        /// <param name="entities"></param>
        public void DeleteInBatch(List<T> entities)
        {
            ApplyOperationInBatch(entities, new ContextOperation(DeleteOperation), true);
        }

        /// <summary>
        /// Deletes all the entities in the Azure Table in batch.
        /// </summary>
        public void DeleteAll()
        {
            this.DeleteInBatch(this.Entities);
        }

        #endregion DELETE

        #endregion CRUD Operations

        /// <summary>
        /// Counts all entities in the Azure Table.
        /// </summary>
        /// <returns>Returns the entities count.</returns>
        public int Count()
        {
            return Count(Entities);
        }

        /// <summary>
        /// Counts all resulting entities of the query execution.
        /// </summary>
        /// <param name="query">Query to execute.</param>
        /// <returns>Returns the entities count.</returns>
        public int Count(IQueryable<T> query)
        {
            var newQuery =
                from e in query
                select new { e.PartitionKey };

            return ExecuteQuery(newQuery).ToList().Count();
        }
        
        /// <summary>
        /// Indicates if an entity with the specified PartitionKey and RowKey exists.
        /// </summary>
        /// <param name="partitionKey">Entity's PartitionKey.</param>
        /// <param name="rowKey">Entity's RowKey.</param>
        /// <returns>Returns a boolean indicating if the entity exists or not.</returns>
        public bool Exists(string partitionKey, string rowKey)
        {
            var query =
                from entity in Entities
                where entity.PartitionKey == partitionKey && entity.RowKey == rowKey
                select entity;

            return Exists(query);
        }

        /// <summary>
        /// Indicates if any resulting entity of the query execution exists.
        /// </summary>
        /// <param name="query">Query to execute.</param>
        /// <returns>Returns a boolean indicating if any entity exists or not.</returns>
        public bool Exists(IQueryable<T> query)
        {
            return Count(query) > 0;
        }

        #region Private
        
        /// <summary>
        /// Internal method used by this context to resolve entity set name from the entity type.
        /// </summary>
        /// <param name="type"></param>
        /// <returns>Returns the entity set name.</returns>
        private string ResolveName(Type type)
        {
            return _entitySetName;
        }

        /// <summary>
        /// Internal method used by this context to resolve the entity type from the entity set name.
        /// </summary>
        /// <param name="entitySetName"></param>
        /// <returns>Returns the entity type.</returns>
        private Type ResolveType(string entitySetName)
        {
            return _entityType;
        }
        
        /// <summary>
        /// Validates an entity.
        /// </summary>
        /// <param name="entity">Entity to validate.</param>
        private void ValidateEntity(T entity)
        {
            if (entity == null)
                throw new Exception(string.Format("No [{0}] entity specified.", _entitySetName));

            string invalidSymbol = "%";
            if (entity.PartitionKey.Contains(invalidSymbol))
                throw new Exception(string.Format("[{0}] entity's Partition Key cannot contain the symbol: {1}", _entitySetName, invalidSymbol));
            if (entity.RowKey.Contains(invalidSymbol))
                throw new Exception(string.Format("[{0}] entity's Row Key cannot contain the symbol: {1}", _entitySetName, invalidSymbol));
        }

        /// <summary>
        /// Apply changes to the Azure Table.
        /// </summary>
        /// <param name="skipConsistency">Specifies whether to skip consistency check when applying changes or not.</param>
        private void SaveChanges(bool skipConsistency)
        {
            SaveChanges(_context.SaveChangesDefaultOptions, skipConsistency);
        }

        /// <summary>
        /// Apply changes to the Azure Table.
        /// </summary>
        /// <param name="options">Save changes options to apply.</param>
        /// <param name="skipExistenceCheck">Specifies whether to skip consistency check when applying changes or not.</param>
        private void SaveChanges(SaveChangesOptions options, bool skipExistenceCheck)
        {
            if (skipExistenceCheck)
            {
                try
                {
                    if (_saveWithRetries)
                        _context.SaveChangesWithRetries();
                    else
                        _context.SaveChanges(options);
                }
                catch (DataServiceRequestException ex)
                {
                    DataServiceClientException innEx = ex.InnerException as DataServiceClientException;
                    if (innEx != null && innEx.StatusCode == 404)
                        return;

                    throw;
                }
                finally
                {
                    RaiseChangesSaved();
                    
                    ClearDeletedEntities();
                }
            }
            else
            {
                if (_saveWithRetries)
                    _context.SaveChangesWithRetries();
                else
                    _context.SaveChanges(options);

                RaiseChangesSaved();

                ClearDeletedEntities();
            }
        }

        /// <summary>
        /// Removes all entities marked as "Deleted" from the internal context.
        /// </summary>
        private void ClearDeletedEntities()
        {
            var deletedEntities = _context.Entities.Where(en => en.State == EntityStates.Deleted);
            foreach (var deletedEntity in deletedEntities)
            {
                _context.Detach(deletedEntity.Entity);
            }
        }

        /// <summary>
        /// Applies the specified operation to the specified entity.
        /// </summary>
        /// <param name="entity">Entity to apply operation to.</param>
        /// <param name="operation">Operation to apply.</param>
        private void ApplyOperation(T entity, ContextOperation operation)
        {
            ApplyOperation(entity, operation, false);
        }

        /// <summary>
        /// Applies the specified operation to the specified entity.
        /// </summary>
        /// <param name="entity">Entity to apply operation to.</param>
        /// <param name="operation">Operation to apply.</param>
        /// <param name="skipExistenceCheck">Specifies whether to skip consistency check when applying the operation.</param>
        private void ApplyOperation(T entity, ContextOperation operation, bool skipExistenceCheck)
        {
            operation(entity, skipExistenceCheck);

            _currentOperationEventArgs.TotalEntities = 1;
            _currentOperationEventArgs.EntitiesSaved = 1;

            SaveChanges(skipExistenceCheck);
        }

        /// <summary>
        /// Applies the specified operation to the specified entities in batch.
        /// </summary>
        /// <param name="entities">Entities to apply operation to.</param>
        /// <param name="operation">Operation to apply.</param>
        private void ApplyOperationInBatch(List<T> entities, ContextOperation operation)
        {
            ApplyOperationInBatch(entities, operation, false);
        }

        /// <summary>
        /// Applies the specified operation to the specified entities in batch.
        /// </summary>
        /// <param name="entities">Entities to apply operation to.</param>
        /// <param name="operation">Operation to apply.</param>
        /// <param name="skipExistenceCheck">Specifies whether to skip consistency check when applying the operation.</param>
        private void ApplyOperationInBatch(List<T> entities, ContextOperation operation, bool skipExistenceCheck)
        {
            if (entities == null)
                return;

            _currentOperationEventArgs.TotalEntities = entities.Count;
            _currentOperationEventArgs.EntitiesSaved = 0;

            var entitiesGroups = entities.GroupBy(en => en.PartitionKey);
            foreach (var entitiesGroup in entitiesGroups)
            {
                string pk = entitiesGroup.Key;

                int counter = 0;
                foreach (T entity in entitiesGroup)
                {
                    operation(entity, skipExistenceCheck);
                    counter++;

                    if (counter == BatchSize)
                    {
                        _currentOperationEventArgs.EntitiesSaved += counter;

                        SaveChanges(SaveChangesOptions.Batch, skipExistenceCheck);
                        counter = 0;
                    }
                }

                if (counter > 0)
                {
                    _currentOperationEventArgs.EntitiesSaved += counter;

                    SaveChanges(SaveChangesOptions.Batch, skipExistenceCheck);
                }
            }
        }

        /// <summary>
        /// Add operation logic to apply to an entity.
        /// </summary>
        /// <param name="entity">Entity to apply operation to.</param>
        /// <param name="skipExistenceCheck">Specifies whether to skip consistency check or not.</param>
        private void AddOperation(T entity, bool skipExistenceCheck)
        {
            _currentOperationEventArgs.OperationName = GenericTableContextOperationName.Add;

            ValidateEntity(entity);

            _context.AddObject(_entitySetName, entity);
        }

        /// <summary>
        /// Update operation logic to apply to an entity.
        /// </summary>
        /// <param name="entity">Entity to apply operation to.</param>
        /// <param name="skipExistenceCheck">Specifies whether to skip consistency check or not.</param>
        private void UpdateOperation(T entity, bool skipExistenceCheck)
        {
            _currentOperationEventArgs.OperationName = GenericTableContextOperationName.Update;

            ValidateEntity(entity);

            T updatedEntity = null;

            if (skipExistenceCheck)
            {
                DetachEntityIfExists(entity);
                _context.AttachTo(_entitySetName, entity, "*");

                updatedEntity = entity;
            }
            else
            {
                T entityInContext = GetEntityInContext(entity, false);
                if (entityInContext != null)
                    updatedEntity = entity;
                else
                {
                    string pk = entity.PartitionKey;
                    string rk = entity.RowKey;

                    updatedEntity = Get(pk, rk);
                    if (updatedEntity == null)
                        throw new Exception(string.Format("The entity [{0}](PartitionKey='{1}', RowKey='{2}') doesn't exist.", _entitySetName, pk, rk));

                    PropertyInfo property = null;
                    int length = _entityProperties.Length;
                    for (int i = length - 1; i >= 0; i--)
                    {
                        property = _entityProperties[i];
                        property.SetValue(updatedEntity, property.GetValue(entity, null), null);
                    }
                }
            }

            _context.UpdateObject(updatedEntity);
        }

        /// <summary>
        /// Delete operation logic to apply to an entity.
        /// </summary>
        /// <param name="entity">Entity to apply operation to.</param>
        /// <param name="skipExistenceCheck">Specifies whether to skip consistency check or not.</param>
        private void DeleteOperation(T entity, bool skipExistenceCheck)
        {
            _currentOperationEventArgs.OperationName = GenericTableContextOperationName.Delete;

            //  Never check for existence when deleting, getting the latest data from the data store is unnecessary
            if (entity == null)
                return;

            DetachEntityIfExists(entity);
            _context.AttachTo(_entitySetName, entity, "*");
            
            _context.DeleteObject(entity);
        }

        /// <summary>
        /// Creates an entity's instance with nothing but the PartitionKey and RowKey.
        /// </summary>
        /// <param name="partitionKey">Entity's PartitionKey.</param>
        /// <param name="rowKey">Entity's RowKey.</param>
        /// <returns>Returns the created entity's instance.</returns>
        private T CreateInstanceOfEntity(string partitionKey, string rowKey)
        {
            T entity = (T)_entityConstructor.Invoke(new object[0] { });
            entity.PartitionKey = partitionKey;
            entity.RowKey = rowKey;
            return entity;
        }

        /// <summary>
        /// Detaches the specified entity from the context.
        /// </summary>
        /// <param name="entity">Entity to detaches.</param>
        private void DetachEntityIfExists(T entity)
        {
            T entityInContext = GetEntityInContext(entity, true);
            if (entityInContext != null)
                _context.Detach(entityInContext);
        }

        /// <summary>
        /// Indicates whether the entity exists in the context or not.
        /// </summary>
        /// <param name="entity">Entity to lookup.</param>
        /// <returns>Returns a boolean whether the entity exists in the context or not.</returns>
        private bool ContextContainsEntity(T entity)
        {
            return GetEntityInContext(entity, false) != null;
        }

        /// <summary>
        /// Gets the specified entity from the context.
        /// </summary>
        /// <param name="entity">Entity to find.</param>
        /// <param name="findByKeys">Specifies whether to find by keys or not.</param>
        /// <returns>Returns the entity from the context.</returns>
        private T GetEntityInContext(T entity, bool findByKeys)
        {
            EntityDescriptor entityDesc = _context.GetEntityDescriptor(entity);

            if (entityDesc != null)
                return entity;

            if (!findByKeys)
                return null;

            foreach (EntityDescriptor entityDesc2 in _context.Entities)
            {
                T entityInContext = (T)entityDesc2.Entity;

                if (entityInContext.PartitionKey == entity.PartitionKey && entityInContext.RowKey == entity.RowKey)
                    return entityInContext;
            }

            return null;
        }

        private void RaiseChangesSaved()
        {
            if (_currentOperationEventArgs == null)
                _currentOperationEventArgs = new GenericTableContextOperationEventArgs();

            OnChangesSaved(_currentOperationEventArgs);
        }

        #endregion Private

        #region Protected

        protected virtual void OnChangesSaved(GenericTableContextOperationEventArgs e)
        {
            if (ChangesSaved != null)
                ChangesSaved(this, e);
        }

        #endregion

        #endregion Methods
    }

    public enum GenericTableContextOperationName
    {
        Add = 0,
        Update = 1,
        Delete = 2
    }

    public class GenericTableContextOperationEventArgs : EventArgs
    {
        /// <summary>
        /// Performed operation to the saved entities.
        /// </summary>
        public GenericTableContextOperationName OperationName { get; set; }
        
        /// <summary>
        /// Number of saved entities.
        /// </summary>
        public int EntitiesSaved { get; set; }
        
        /// <summary>
        /// Number of total entities.
        /// </summary>
        public int TotalEntities { get; set; }
    }
}
